import * as vscode from "vscode";

export class MermaidLinter {
  private diagnostics: vscode.DiagnosticCollection;
  private astCache = new Map<string, any>();
  private mermaid: any;
  private initialized: boolean = false;

  constructor() {
    this.diagnostics = vscode.languages.createDiagnosticCollection("mermaid");
    this.initialize();
  }

  private async initialize() {
    if (!this.initialized) {
      const mermaid = await import("mermaid");
      this.mermaid = mermaid.default;
      this.initialized = true;
    }
  }

  async check(document: vscode.TextDocument): Promise<vscode.Diagnostic[]> {
    // 确保mermaid已初始化
    if (!this.initialized) {
      await this.initialize();
    }

    const text = document.getText();
    const errors: vscode.Diagnostic[] = [];

    try {
      // 使用AST缓存
      const ast = await this.parseWithCache(text);
      errors.push(...this.checkCustomRules(ast, document));
    } catch (err) {
      errors.push(this.parseMermaidError(err as Error, document));
    }

    this.diagnostics.set(document.uri, errors);
    return errors;
  }

  private async parseWithCache(code: string): Promise<any> {
    const hash = this.createCodeHash(code);
    if (this.astCache.has(hash)) {
      return this.astCache.get(hash)!;
    }

    const ast = await this.mermaid.parse(code);
    this.astCache.set(hash, ast);
    return ast;
  }

  private checkCustomRules(
    ast: any,
    doc: vscode.TextDocument
  ): vscode.Diagnostic[] {
    const errors: vscode.Diagnostic[] = [];

    // 示例规则：检测孤立节点
    if (ast.type === "flowchart") {
      const connectedNodes = new Set<string>();

      ast.edges.forEach((edge: any) => {
        connectedNodes.add(edge.start);
        connectedNodes.add(edge.end);
      });

      ast.nodes.forEach((node: any) => {
        if (!connectedNodes.has(node.id)) {
          errors.push({
            severity: vscode.DiagnosticSeverity.Warning,
            range: this.findNodeRange(node, doc),
            message: `孤立节点 "${node.id}"`,
            source: "mermaid-linter",
            code: "unconnected-node",
          });
        }
      });
    }

    return errors;
  }

  private parseMermaidError(
    err: Error,
    doc: vscode.TextDocument
  ): vscode.Diagnostic {
    const line = this.parseErrorLine(err.message, doc);
    return {
      severity: vscode.DiagnosticSeverity.Error,
      range: new vscode.Range(line, 0, line, Number.MAX_VALUE),
      message: this.simplifyErrorMessage(err.message),
      source: "mermaid",
    };
  }

  private findNodeRange(node: any, doc: vscode.TextDocument): vscode.Range {
    // 根据节点ID在文档中查找实际位置（需要更精确的实现）
    const text = doc.getText();
    const match = new RegExp(`\\b${node.id}\\b`).exec(text);
    if (match) {
      const pos = doc.positionAt(match.index!);
      return new vscode.Range(
        pos,
        new vscode.Position(pos.line, pos.character + node.id.length)
      );
    }
    return new vscode.Range(0, 0, 0, 0);
  }

  private createCodeHash(code: string): string {
    // 更高效的哈希算法
    let hash = 0;
    for (let i = 0; i < code.length; i++) {
      hash = (Math.imul(31, hash) + code.charCodeAt(i)) | 0;
    }
    return "ast-" + hash.toString(36);
  }

  private simplifyErrorMessage(msg: string): string {
    return msg.replace(/^.*?:\s*/, "").replace(/\s*\(\d+:\d+\).*$/, "");
  }

  private parseErrorLine(message: string, doc: vscode.TextDocument): number {
    const lineMatch = message.match(/line (\d+)/);
    return lineMatch
      ? Math.min(Number(lineMatch[1]) - 1, doc.lineCount - 1)
      : 0;
  }
}
